#========================================
#    author: Changlong.Zang
#      mail: zclongpop123@163.com
#      time: Tue Apr 18 10:49:29 2017
#========================================
try:
    from PyQt5 import QtCore, QtGui
except:
    from PyQt4 import QtCore, QtGui

import maya.cmds as mc
import stCore
#--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
class ListModel(QtCore.QAbstractListModel):
    '''
    '''
    def __init__(self, parent=None):
        '''
        '''
        super(ListModel, self).__init__(parent)
        self.__data = list()





    def rowCount(self, index=QtCore.QModelIndex()):
        '''
        '''
        return len(self.__data)





    def data(self, index=QtCore.QModelIndex(), role=QtCore.Qt.DisplayRole):
        '''
        '''
        if role == QtCore.Qt.DisplayRole:
            return self.__data[index.row()]





    def flags(self, index=QtCore.QModelIndex()):
        '''
        '''
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable





    def insertRow(self, row, data, index=QtCore.QModelIndex()):
        '''
        '''
        realRow = min(row, len(self.__data))
        self.beginInsertRows(index, realRow, realRow)
        self.__data.insert(realRow, data)
        self.endInsertRows()





    def removeRow(self, row, index=QtCore.QModelIndex()):
        '''
        '''
        self.beginRemoveRows(index, row, row)
        if row < len(self.__data):
            self.__data.pop(row)
        self.endRemoveRows()





    def updateData(self, newData=list()):
        '''
        '''
        for i in reversed(range(self.rowCount())):
            self.removeRow(i)

        for dt in newData:
            self.insertRow(self.rowCount(), dt)





class TableModel(QtCore.QAbstractTableModel, QtGui.QWidget):
    '''
    '''
    __font = QtGui.QFont()
    __font.setPointSize(12)
    __font.setBold(True)
    __anim_start = int(mc.playbackOptions(q=True, ast=True))
    __anim_end   = int(mc.playbackOptions(q=True, aet=True)) + 1
    __anim_fps   = dict((('film', 24), ('pal', 25), ('nstc', 30), ('game', 15))).get(mc.currentUnit(q=True, time=True))

    def __init__(self, parent=None):
        '''
        '''
        super(TableModel, self).__init__(parent)
        self.__data        = dict()
        self.__header_data = list()





    def columnCount(self, index=QtCore.QModelIndex()):
        '''
        '''
        return len(self.__data)





    def rowCount(self, index=QtCore.QModelIndex()):
        '''
        '''
        return (self.__anim_end - self.__anim_start)





    def data(self, index=QtCore.QModelIndex(), role=QtCore.Qt.DisplayRole):
        '''
        '''
        if role == QtCore.Qt.DisplayRole:
            key = self.__header_data[index.column()]
            return self.__data[key][index.row()]

        if role == QtCore.Qt.TextAlignmentRole:
            return QtCore.Qt.AlignHCenter | QtCore.Qt.AlignVCenter

        if role == QtCore.Qt.FontRole:
            return self.__font

        if role == QtCore.Qt.ForegroundRole:
            if self.__header_data[index.column()].endswith('Orig'):
                return QtGui.QColor(0, 255, 0)
            else:
                return QtGui.QColor(255, 0, 0)





    def flags(self, index=QtCore.QModelIndex()):
        '''
        '''
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsDragEnabled | QtCore.Qt.ItemIsDropEnabled





    def headerData(self, sec, orient, role=QtCore.Qt.DisplayRole):
        '''
        '''
        if role == QtCore.Qt.DisplayRole:
            if orient == QtCore.Qt.Vertical:
                return '{0:d} - {1:d}'.format(int(sec / self.__anim_fps) + 1, sec + 1)

            else:
                return self.__header_data[sec]





    def supportedDragActions(self):
        return QtCore.Qt.MoveAction





    def supportedDropActions(self):
        return QtCore.Qt.MoveAction





    def dropMimeData(self, data, action, row, column, parent):
        '''
        '''
        self._src_namespace = self.headerData(self._src_index.column(), QtCore.Qt.Horizontal)
        self._src_frame     = self.__anim_start + self._src_index.row()

        self._dst_namespace = self.headerData(parent.column(), QtCore.Qt.Horizontal)
        self._dst_frame     = self.__anim_start + parent.row()

        stCore.move_keys(self._src_namespace, self._src_frame, self._dst_namespace, self._dst_frame)
        return True





    def updateData(self, data):
        '''
        '''
        #- Clear Old Data
        while self.columnCount():
            self.beginRemoveColumns(QtCore.QModelIndex(), 0, 0)
            self.__data.pop(sorted(self.__data.keys())[0])
            self.endRemoveColumns()

        del self.__header_data[:]

        #- Rebuild row count...
        while self.rowCount():
            self.beginRemoveRows(QtCore.QModelIndex(), 0, self.rowCount())
            self.__anim_end -= 1
            self.endRemoveRows()

        self.__anim_start = int(mc.playbackOptions(q=True, ast=True))
        self.__anim_end   = int(mc.playbackOptions(q=True, ast=True))
        new_end           = int(mc.playbackOptions(q=True, aet=True)) + 1
        for i in range(self.__anim_start, new_end):
            self.beginInsertRows(QtCore.QModelIndex(), self.__anim_end, self.__anim_end)
            self.__anim_end += 1
            self.endInsertRows()

        #- Add New Data
        for k, v in data.iteritems():
            self.beginInsertColumns(QtCore.QModelIndex(), self.columnCount(), self.columnCount())

            keylist = list()
            x = 1
            for i in range(self.__anim_start, self.__anim_end):
                if i in v:
                    keylist.append(x)
                    x += 1
                else:
                    keylist.append(None)

            self.__data[k] = keylist
            self.__header_data.append(k)
            self.__header_data.sort()

            self.endInsertColumns()
